import torch

import torchvision.datasets as dataset
from torchvision import transforms

train_data = dataset.MNIST(
    root="mnist", train=True, transform=transforms.ToTensor(), download=False)
#每批64个图片，打乱顺序
train_loader = torch.utils.data.DataLoader(
    train_data, batch_size=64, shuffle=True)

class CNN(torch.nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(1, 32, kernel_size=5, stride=1, padding=2),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2)
        )
        self.fc = torch.nn.Linear(14*14*32, 10)
    def forward(self, x):
        out = self.conv1(x)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out


accuracy = 0
full = 0
model = CNN()
model.load_state_dict(torch.load("mnist.model"))
model.eval()
for (images, labels) in train_loader:
    #验证的图片是28*28的灰度图像
    outputs = model(images)
    _, predicted = outputs.max(1)
    accuracy += (predicted == labels).sum().item()
    full += len(labels)

print(accuracy / full)
